# encoding:utf-8
"""
@project:Neo4j Data ETL
@author:liuzeyu
@time:2022-08-09
@version:v1.1
@Update description:
v1.1（2022-08-10）更新说明
（1）修改了股票节点的生成逻辑，新增了“股票名称”属性
（2）支持对股票&基金关系进行剪枝，开启“prune”参数可以过滤掉持仓不超过1%的股票
（3）可以在程序的入口函数中通过设置布尔值参数来选择输出哪些节点或关系
v1.0（2022-08-08）更新说明
（1）修改配置文件键值对存储格式，节点使用node前缀、关系使用relation前缀、数据使用Data前缀，使之易于区分
（2）新增Neo4jManager和Neo4jIndustry两个类：
    Neo4jManager生成manager节点CSV文件，以及manager_fund关系csv文件
    Neo4jIndustry生成industry节点CSV文件，以及industry_stock关系csv文件
"""
import pandas as pd
import configparser as cp
from memory_profiler import profile


class Neo4j:
    def __init__(self, neo4j_config: cp.ConfigParser):
        self.config = neo4j_config
        self.file_path = neo4j_config.get('data_path', 'file_path')
        self.file_path_stockname = neo4j_config.get('data_path', 'file_path_stockname')
        self.stock_node_path = neo4j_config.get('node_export_path', 'stock_node')
        self.fund_node_path = neo4j_config.get('node_export_path', 'fund_node')
        self.relation_path = neo4j_config.get('relation_export_path', 'stock_fund_relation')
        self.language = neo4j_config.get('env', 'language')
        self.export = neo4j_config.get('operation', 'export')
        self.prune = neo4j_config.get('operation', 'prune')
        if self.prune:
            self.prune_threshold = float(neo4j_config.get('operation', 'prune_threshold'))

    # def rename_columns(self):
    #     """
    #
    #     :return:
    #     """
    #     df = pd.read_csv(self.file_path)
    #     if self.language == 'EN':
    #         df.rename(columns={
    #             '公告日': '',
    #             '基金代码': 'fund_code',
    #             '股票代码': 'stock_code',
    #             '报告期': '',
    #             '占流通股本比例（%）': '',
    #             '持有股票数量（股）': '',
    #             '持有股票市值（元）': '',
    #             '持有股票市值占基金净值比例（%）': '',
    #             '管理人': 'fund_manager',
    #             '名称': 'fund_name'
    #         }, inplace=True)
    #     if self.export:
    #         df.to_csv('./Data_EN.csv', index='index')
    #     return df

    @property
    def stock_node(self):
        """
        从源数据中抽取
        :return:
        """
        df = pd.read_csv(self.file_path, low_memory=False, dtype=str).loc[:, ['股票代码']]
        name = pd.read_csv(self.file_path_stockname, dtype=str)
        stock_node = pd.DataFrame(df).sort_values(by=['股票代码']).drop_duplicates().reset_index(drop=True)
        stock_node = pd.merge(stock_node, name, how='outer', on=['股票代码', '股票代码'])
        if self.language == 'EN':
            stock_node.rename(columns={'股票代码': 'stock_code', '股票简称': 'stock_name'}, inplace=True)
        if self.export:
            stock_node.to_csv(self.stock_node_path, index_label='s_id')
        return stock_node

    @property
    def fund_node(self):
        """

        :return:
        """
        df = pd.read_csv(self.file_path, low_memory=False).loc[:, ['基金代码', '管理人', '名称']]
        fund_node = pd.DataFrame(df).sort_values(by=['基金代码']).drop_duplicates().reset_index(drop=True)
        if self.language == 'EN':
            fund_node.rename(columns={'基金代码': 'fund_code', '管理人': 'fund_manager', '名称': 'fund_name'}, inplace=True)
        if self.export:
            fund_node.to_csv(self.fund_node_path, index_label='f_id')
        return fund_node

    @property
    def stock_fund_relation(self):
        """

        :return:
        """
        df = pd.read_csv(self.file_path, low_memory=False).loc[:, ['基金代码', '股票代码', '持有股票市值占基金净值比例（%）']]
        relation = pd.DataFrame(df).drop_duplicates().reset_index(drop=True)
        if self.language == 'EN':
            mapper_f = pd.read_csv(self.fund_node_path, index_col=False, dtype=str).loc[:, ['fund_code', 'f_id']].set_index('fund_code').to_dict()
            mapper_s = pd.read_csv(self.stock_node_path, index_col=False, dtype=str).loc[:, ['stock_code', 's_id']].set_index('stock_code').to_dict()
            relation['基金代码'] = relation['基金代码'].map(mapper_f['f_id'])
            relation['股票代码'] = relation['股票代码'].map(mapper_s['s_id'])
            relation.rename(columns={'基金代码': 'f_id', '股票代码': 's_id', '持有股票市值占基金净值比例（%）': 'Ratio_hold'}, inplace=True)
        else:
            mapper_f = pd.read_csv(self.fund_node_path, index_col=False, dtype=str).loc[:, ['基金代码', 'f_id']].set_index('基金代码').to_dict()
            mapper_s = pd.read_csv(self.stock_node_path, index_col=False, dtype=str).loc[:, ['股票代码', 's_id']].set_index('股票代码').to_dict()
            relation['基金代码'] = relation['基金代码'].map(mapper_f['f_id'])
            relation['股票代码'] = relation['股票代码'].map(mapper_s['s_id'])
            relation.rename(columns={'基金代码': 'f_id', '股票代码': 's_id', '持有股票市值占基金净值比例（%）': '持有股票市值占基金净值比例'}, inplace=True)
        if self.prune:
            relation['Ratio_hold'] = pd.to_numeric(relation['Ratio_hold'], errors='coerce')
            relation = relation[relation['Ratio_hold'] > self.prune_threshold]
        if self.export:
            relation.to_csv(self.relation_path, index_label='r_id')
        return relation


class Neo4jManager:
    def __init__(self, neo4j_config: cp.ConfigParser):
        self.config = neo4j_config
        self.file_path = neo4j_config.get('data_path', 'file_path_manager')
        self.manager_node_path = neo4j_config.get('node_export_path', 'manager_node')
        self.fund_node_path = neo4j_config.get('node_export_path', 'fund_node')
        self.relation_path = neo4j_config.get('relation_export_path', 'manager_fund_relation')
        self.language = neo4j_config.get('env', 'language')
        self.drop = self.export = neo4j_config.get('operation', 'drop')
        self.export = self.export = neo4j_config.get('operation', 'export')

    @property
    def manager_node(self):
        """

        :return:
        """
        manager_node = pd.read_csv(self.file_path, dtype=str).loc[:, ['基金经理姓名']]
        manager_node.drop_duplicates(subset='基金经理姓名', keep='first', inplace=True)
        manager_node = manager_node.reset_index(drop=True)
        if self.language == 'EN':
            manager_node.rename(columns={'基金经理姓名': 'name'}, inplace=True)
        if self.export:
            manager_node.to_csv(self.manager_node_path, index_label='m_id')
        return manager_node

    @property
    def manager_fund_relation(self):
        """

        :return:
        """
        relation = pd.read_csv(self.file_path).loc[:, ['基金经理姓名', '基金代码', '任职日期', '离职日期']]
        mapper_f = pd.read_csv(self.fund_node_path, index_col=False, dtype=str).loc[:, ['fund_code', 'f_id']].set_index('fund_code').to_dict()
        mapper_m = pd.read_csv(self.manager_node_path, index_col=False, dtype=str).loc[:, ['name', 'm_id']].set_index('name').to_dict()
        relation['基金代码'] = relation['基金代码'].map(mapper_f['f_id'])
        relation['基金经理姓名'] = relation['基金经理姓名'].map(mapper_m['m_id'])
        relation.rename(columns={'基金代码': 'f_id', '基金经理姓名': 'm_id', '任职日期': 'entrance_date', '离职日期': 'departure_date'}, inplace=True)
        if self.drop:
            relation = relation.dropna(subset='f_id')
        if self.export:
            relation.to_csv(self.relation_path, index_label='r_id')
        return relation


class Neo4jIndustry:
    def __init__(self, neo4j_config: cp.ConfigParser):
        self.config = neo4j_config
        self.file_path = neo4j_config.get('data_path', 'file_path_industry')
        self.file_path_r = neo4j_config.get('data_path', 'file_path_industry_r')
        self.industry_node_path = neo4j_config.get('node_export_path', 'industry_node')
        self.stock_node_path = neo4j_config.get('node_export_path', 'stock_node')
        self.relation_path = neo4j_config.get('relation_export_path', 'industry_stock_relation')
        self.language = neo4j_config.get('env', 'language')
        self.drop = self.export = neo4j_config.get('operation', 'drop')
        self.export = self.export = neo4j_config.get('operation', 'export')

    @property
    def industry_node(self):
        """

        :return:
        """
        industry_node = pd.read_csv(self.file_path, dtype=str).sort_values(by=['行业代码'])
        industry_node = industry_node.reset_index(drop=True)
        if self.language == 'EN':
            industry_node.rename(columns={'行业代码': 'industry_code', '申万行业': 'SYWG_industry'}, inplace=True)
        if self.export:
            industry_node.to_csv(self.industry_node_path, index_label='i_id')
        return industry_node

    @property
    def industry_stock_relation(self):
        """

        :return:
        """
        relation = pd.read_csv(self.file_path_r)
        mapper_f = pd.read_csv(self.stock_node_path, index_col=False, dtype=str).loc[:, ['stock_code', 's_id']].set_index('stock_code').to_dict()
        mapper_m = pd.read_csv(self.industry_node_path, index_col=False, dtype=str).loc[:, ['industry_code', 'i_id']].set_index('industry_code').to_dict()
        relation['行业代码'] = relation['行业代码'].astype(str).str[:8]
        relation['股票代码'] = relation['股票代码'].map(mapper_f['s_id'])
        relation['行业代码'] = relation['行业代码'].map(mapper_m['i_id'])
        relation.rename(columns={'股票代码': 's_id', '行业代码': 'i_id'}, inplace=True)
        if self.drop:
            relation = relation.dropna(subset='s_id')
        if self.export:
            relation.to_csv(self.relation_path, index_label='r_id')
        return relation


# def cut_csv(df: pd.DataFrame) -> None:
#     """
#
#     :param df:
#     :return:
#     """
#     for i in range(5):
#         dset = df.iloc[i:i * 100000, :]
#         f_name = './relation_edge_part' + str(i) + '.csv'
#         dset.to_csv(f_name, index='r_id')

@profile
def main(node_stock: bool = False, node_fund: bool = False, node_manager: bool = False, node_industry: bool = False, relation_stock_fund: bool = False, relation_manager_fund: bool = False, relation_industry_stock: bool = False):
    # 加载配置文件
    config = cp.ConfigParser()
    config.read('config.ini', encoding='utf-8')

    # 生成stock、fund两个节点CSV文件，以及stock_fund关系csv文件
    N = Neo4j(neo4j_config=config)
    if node_stock:
        stock = N.stock_node
    if node_fund:
        fund = N.fund_node
    if relation_stock_fund:
        stock_fund_relation = N.stock_fund_relation

    # 生成manager节点CSV文件，以及manager_fund关系csv文件
    M = Neo4jManager(neo4j_config=config)
    if node_manager:
        manager = M.manager_node
    if relation_manager_fund:
        manager_fund_relation = M.manager_fund_relation

    # 生成industry节点CSV文件，以及industry_stock关系csv文件
    I = Neo4jIndustry(neo4j_config=config)
    if node_industry:
        industry = I.industry_node
    if relation_industry_stock:
        industry_stock_relation = I.industry_stock_relation

    print('ETL过程执行完毕')


if __name__ == '__main__':
    main(
        node_stock=True, node_fund=True, node_manager=True, node_industry=True,
        relation_stock_fund=True, relation_manager_fund=True, relation_industry_stock=True
    )